################################################################################
#
# Useful function
#
################################################################################

import pyflag.DB as DB
import pyflag.FlagFramework as FlagFramework
import pyflag.pyflaglog as pyflaglog
import pyflag.conf

config = pyflag.conf.ConfObject()
config.add_option("NSRL_HASH_COUNT", default = "HashCount")
config.add_option("NSRL_PRODUCT_COUNT", default = "ProductCount")

def check_index():
    ## Check if ext nsrl tables exists
    pyflaglog.log(pyflaglog.DEBUG, "Checking NSRL tables...")
    try:
        dbh = DB.DBO("information_schema")
        dbh.execute("""SELECT * FROM TABLES
            WHERE TABLE_SCHEMA = '%s' AND (TABLE_NAME = '%s' OR TABLE_NAME = '%s')""",
            config.NSRL_DB, config.NSRL_HASH_COUNT, config.NSRL_PRODUCT_COUNT)
        row = dbh.fetch()

        if row == None:
            dbh = DB.DBO(config.NSRL_DB)
            init_index(dbh)
    except Exception, e:
        pyflaglog.log(pyflaglog.ERROR, "Can not create ext nsrl tables")
	pyflaglog.log(pyflaglog.ERROR, str(e))


def init_index(dbh):
    pyflaglog.log(pyflaglog.DEBUG, "Creating ext nsrl tables...")
    pyflaglog.log(pyflaglog.INFO, "This may take several minutes...")

    dbh.execute("""
        INSERT INTO `%s` (
	SELECT SHA1, MD5, count(*) AS Count
        FROM (SELECT DISTINCT SHA1, MD5, ProductCode FROM `%s`) A
        GROUP BY SHA1, MD5)""",
        (config.NSRL_HASH_COUNT, config.NSRL_HASHES))

    dbh.execute("""
	INSERT INTO `%s` (
	SELECT C.ProductCode,
        count(*) AS NSRLCount,
        min(C.N) AS NSRLMin,
        sum(1 / C.N) AS NSRLNormalSum,
        sum(power(1 / C.N, 2)) AS NSRLExtraSum,
        sum(power(1 / C.N, 3)) AS NSRLUltraSum,
	sum((sign((16 - C.N) / 15) + 1) * (16 - C.N) / 30) AS LIN_15_sum,
        sum((sign((11 - C.N) / 10) + 1) * (11 - C.N) / 20) AS LIN_10_sum,
        sum((sign((6 - C.N) / 5) + 1) * (6 - C.N) / 10) AS LIN_5_sum
        FROM (SELECT A.ProductCode, B.Count AS N FROM (SELECT DISTINCT SHA1, MD5, ProductCode FROM `%s`) A 
	INNER JOIN `%s` B ON A.SHA1 = B.SHA1 AND A.MD5 = B.MD5) C
        GROUP BY C.ProductCode)""",
        (config.NSRL_PRODUCT_COUNT, config.NSRL_HASHES, config.NSRL_HASH_COUNT))



## Get ProductCode list for SHA1 & MD5 hash
def getProductCodeList(sha1, md5):
    productCodeList = []
    dbh = DB.DBO(config.NSRL_DB)
    dbh.execute("""SELECT DISTINCT CONVERT(ProductCode, CHAR(8)) AS ProductCode FROM %s WHERE SHA1 = '%s' AND MD5 = '%s'""", 
	(config.NSRL_HASHES, sha1, md5))
    row = dbh.fetch()

    while row:
        productCode = row['ProductCode']
        productCodeList.append(productCode)
        row = dbh.fetch()

    return productCodeList
										


## Get Membership value for SHA1 & MD5
def getMembership(sha1, md5):
    membership = 0.0
    dbh = DB.DBO(config.NSRL_DB)
    dbh.execute("""SELECT 1 / Count AS Membership FROM %s WHERE SHA1 = '%s' AND MD5 = '%s'""", 
        (config.NSRL_HASH_COUNT, sha1, md5))
    row = dbh.fetch()

    if row != None: membership = row['Membership']

    return membership



## Get Count for SHA1 & MD5
def getHashCount(sha1, md5):
    count = 0
    dbh = DB.DBO(config.NSRL_DB)
    dbh.execute("""SELECT Count FROM %s WHERE SHA1= '%s' AND MD5 = '%s'""", 
        (config.NSRL_HASH_COUNT, sha1, md5))
    row = dbh.fetch()

    if row != None: count = row['Count']

    return count



## Check if hash is into NSRL db
def getProductCount(sha1, md5):
    count = 0
    dbh = DB.DBO(config.NSRL_DB)
    dbh.execute("""SELECT Count FROM %s WHERE SHA1 = '%s' AND MD5 = '%s'""", config.NSRL_HASH_COUNT, sha1, md5)
    row = dbh.fetch()

    if row != None: count = row['Count']

    return count


## Get product set in sql format
from string import join

def getSQLProductSet(productCodeList):
    result = "("
    result += join(productCodeList, ',')
    result += ")"
    return result
